//
// Created by chizuru on 2025/10/15.
//
#pragma once

#include <memory>
#include <functional>
#include <unordered_map>
#include <shared_mutex>
#include <mutex>
#include <string>
#include <typeindex>
#include <stdexcept>
#include <iostream>
#include <atomic>
#include <thread>
#include <vector>
#include <algorithm>

namespace registry {

    class RegistryException : public std::runtime_error {
    public:
        using std::runtime_error::runtime_error;
    };

    class RegistryContainer {
    public:
        static RegistryContainer &instance() {
            static RegistryContainer inst;
            return inst;
        }

        RegistryContainer(const RegistryContainer &) = delete;

        RegistryContainer &operator=(const RegistryContainer &) = delete;

        RegistryContainer(RegistryContainer &&) = delete;

        RegistryContainer &operator=(RegistryContainer &&) = delete;

        ~RegistryContainer() {
            std::vector<std::pair<Key, EntryPtr>> entries;
            {
                std::unique_lock lock(_mapMutex);
                for (auto &p: _map) entries.emplace_back(p.first, p.second);
                _map.clear();
            }
            std::sort(entries.begin(), entries.end(), [](auto const &a, auto const &b) {
                return a.second->destroyPriority > b.second->destroyPriority; // higher first
            });

            for (auto &kv: entries) {
                auto &entry = kv.second;
                std::shared_ptr<void> inst = std::atomic_load(&entry->instance);
                if (!inst) {
                    std::lock_guard guard(entry->entryMutex);
                    inst = entry->weakInstance.lock();
                }
                if (inst) {
                    if (entry->onDestroy) entry->onDestroy(inst);
                }
                entry->factory = nullptr;
                std::atomic_store(&entry->instance, std::shared_ptr<void>());
                std::lock_guard guard(entry->entryMutex);
                entry->weakInstance.reset();
            }
        }

        template<typename T, typename Factory>
        void registerFactory(const std::string &name,
                             Factory &&factory,
                             bool lazy = true,
                             bool useWeakCache = false,
                             int destroyPriority = 0,
                             std::function<void(std::shared_ptr<T>)> onCreate = nullptr,
                             std::function<void(std::shared_ptr<T>)> onDestroy = nullptr) {
            auto key = makeKey<T>(name);

            auto wrappedFactory = [f = std::forward<Factory>(factory)]() -> std::shared_ptr<void> {
                auto p = f();
                if (!p) return nullptr;
                return std::static_pointer_cast<void>(p);
            };

            std::function < void(std::shared_ptr<void>) > wrappedOnCreate = nullptr;
            std::function < void(std::shared_ptr<void>) > wrappedOnDestroy = nullptr;
            if (onCreate) {
                wrappedOnCreate = [onCreate = std::move(onCreate)](std::shared_ptr<void> p) {
                    onCreate(std::static_pointer_cast<T>(p));
                };
            }
            if (onDestroy) {
                wrappedOnDestroy = [onDestroy = std::move(onDestroy)](std::shared_ptr<void> p) {
                    onDestroy(std::static_pointer_cast<T>(p));
                };
            }

            EntryPtr entry = std::make_shared<Entry>();
            entry->factory = std::move(wrappedFactory);
            entry->lazy = lazy;
            entry->useWeakCache = useWeakCache;
            entry->destroyPriority = destroyPriority;
            entry->onCreate = std::move(wrappedOnCreate);
            entry->onDestroy = std::move(wrappedOnDestroy);

            {
                std::unique_lock lock(_mapMutex);
                if (_map.find(key) != _map.end()) {
                    throw RegistryException("duplicate registration for same type+name");
                }
                _map.emplace(std::move(key), entry);
            }

            if (!lazy) {
                std::shared_ptr<void> inst = entry->factory();
                if (!inst) throw RegistryException("factory returned nullptr on eager registration");
                if (entry->useWeakCache) {
                    std::lock_guard guard(entry->entryMutex);
                    entry->weakInstance = inst;
                } else {
                    std::atomic_store(&entry->instance, inst);
                }
                if (entry->onCreate) entry->onCreate(inst);
                entry->factory = nullptr;
            }
        }

        template<typename T>
        void registerInstance(const std::string &name,
                              std::shared_ptr<T> instance,
                              int destroyPriority = 0,
                              std::function<void(std::shared_ptr<T>)> onDestroy = nullptr,
                              bool useWeakCache = false) {
            if (!instance) throw RegistryException("cannot register nullptr instance");
            auto key = makeKey<T>(name);

            EntryPtr entry = std::make_shared<Entry>();
            entry->lazy = false;
            entry->useWeakCache = useWeakCache;
            entry->destroyPriority = destroyPriority;
            if (onDestroy) {
                entry->onDestroy = [onDestroy = std::move(onDestroy)](std::shared_ptr<void> p) {
                    onDestroy(std::static_pointer_cast<T>(p));
                };
            }

            if (useWeakCache) {
                std::lock_guard guard(entry->entryMutex);
                entry->weakInstance = std::static_pointer_cast<void>(instance);
            } else {
                std::atomic_store(&entry->instance, std::static_pointer_cast<void>(instance));
            }

            std::unique_lock lock(_mapMutex);
            if (_map.find(key) != _map.end()) {
                throw RegistryException("duplicate registration for same type+name");
            }
            _map.emplace(std::move(key), entry);
        }

        template<typename T>
        std::shared_ptr<T> get(const std::string &name) {
            auto key = makeKey<T>(name);
            EntryPtr entry;
            {
                std::shared_lock lock(_mapMutex);
                auto it = _map.find(key);
                if (it == _map.end()) {
                    throw RegistryException("no registered service for given type+name");
                }
                entry = it->second;
            }

            std::shared_ptr<void> inst = std::atomic_load(&entry->instance);
            if (inst) return std::static_pointer_cast<T>(inst);

            if (entry->useWeakCache) {
                std::lock_guard guard(entry->entryMutex);
                auto locked = entry->weakInstance.lock();
                if (locked) return std::static_pointer_cast<T>(locked);
            }

            std::call_once(entry->initFlag, [entry]() {
                if (!entry->factory) throw RegistryException("no factory for lazy-initialized entry");
                std::shared_ptr<void> created = entry->factory();
                if (!created) throw RegistryException("factory returned nullptr during lazy initialization");

                if (entry->useWeakCache) {
                    std::lock_guard guard(entry->entryMutex);
                    entry->weakInstance = created;
                } else {
                    std::atomic_store(&entry->instance, created);
                }

                if (entry->onCreate) entry->onCreate(created);
                entry->factory = nullptr;
            });

            inst = std::atomic_load(&entry->instance);
            if (inst) return std::static_pointer_cast<T>(inst);
            std::lock_guard guard(entry->entryMutex);
            auto locked = entry->weakInstance.lock();
            if (locked) return std::static_pointer_cast<T>(locked);

            throw RegistryException("failed to initialize instance");
        }

        template<typename T>
        bool contains(const std::string &name) {
            auto key = makeKey<T>(name);
            std::shared_lock lock(_mapMutex);
            return _map.find(key) != _map.end();
        }

        template<typename T>
        void unregisterService(const std::string &name, bool callDestroy = true) {
            auto key = makeKey<T>(name);
            EntryPtr entry;
            {
                std::unique_lock lock(_mapMutex);
                auto it = _map.find(key);
                if (it == _map.end()) return;
                entry = it->second;
                _map.erase(it);
            }

            std::shared_ptr<void> inst = std::atomic_load(&entry->instance);
            if (!inst) {
                std::lock_guard guard(entry->entryMutex);
                inst = entry->weakInstance.lock();
            }
            if (inst && callDestroy && entry->onDestroy) entry->onDestroy(inst);

            entry->factory = nullptr;
            std::atomic_store(&entry->instance, std::shared_ptr<void>());
            std::lock_guard guard(entry->entryMutex);
            entry->weakInstance.reset();
        }

    private:
        RegistryContainer() = default;

        struct Entry {
            std::function<std::shared_ptr<void>()> factory;
            std::shared_ptr<void> instance;
            std::weak_ptr<void> weakInstance;
            std::once_flag initFlag;
            bool lazy = true;
            bool useWeakCache = false;
            int destroyPriority = 0;
            std::mutex entryMutex;
            std::function<void(std::shared_ptr<void>)> onCreate;
            std::function<void(std::shared_ptr<void>)> onDestroy;
        };

        using EntryPtr = std::shared_ptr<Entry>;

        struct Key {
            std::type_index type;
            std::string name;
        };

        struct KeyHash {
            size_t operator()(Key const &k) const noexcept {
                size_t h1 = std::hash<std::type_index>{}(k.type);
                size_t h2 = std::hash<std::string>{}(k.name);
                return h1 ^ (h2 + 0x9e3779b97f4a7c15ULL + (h1 << 6) + (h1 >> 2));
            }
        };

        struct KeyEq {
            bool operator()(Key const &a, Key const &b) const noexcept {
                return a.type == b.type && a.name == b.name;
            }
        };

        template<typename T>
        static Key makeKey(const std::string &name) {
            return Key{std::type_index(typeid(T)), name};
        }

        std::unordered_map<Key, EntryPtr, KeyHash, KeyEq> _map;
        mutable std::shared_mutex _mapMutex;
    };

    template<typename T, typename Factory>
    struct AutoFactoryRegistrar {
        AutoFactoryRegistrar(const char *name,
                             Factory &&factory,
                             bool lazy = true,
                             bool useWeakCache = false,
                             int destroyPriority = 0,
                             std::function<void(std::shared_ptr<T>)> onCreate = nullptr,
                             std::function<void(std::shared_ptr<T>)> onDestroy = nullptr) {
            RegistryContainer::instance().registerFactory<T>(
                    std::string(name),
                    std::forward<Factory>(factory),
                    lazy,
                    useWeakCache,
                    destroyPriority,
                    std::move(onCreate),
                    std::move(onDestroy)
            );
        }
    };

    template<typename T>
    struct AutoInstanceRegistrar {
        AutoInstanceRegistrar(const char *name,
                              std::shared_ptr<T> instance,
                              int destroyPriority = 0,
                              std::function<void(std::shared_ptr<T>)> onDestroy = nullptr,
                              bool useWeakCache = false) {
            RegistryContainer::instance().registerInstance<T>(
                    std::string(name),
                    std::move(instance),
                    destroyPriority,
                    std::move(onDestroy),
                    useWeakCache
            );
        }
    };

    template<typename T, typename Factory>
    inline auto makeAutoFactoryRegistrar(const char *name, Factory &&f,
                                         bool lazy = true,
                                         bool useWeakCache = false,
                                         int destroyPriority = 0,
                                         std::function<void(std::shared_ptr<T>)> onCreate = nullptr,
                                         std::function<void(std::shared_ptr<T>)> onDestroy = nullptr) {
        return AutoFactoryRegistrar<T, std::decay_t<Factory>>(
                name, std::forward<Factory>(f), lazy, useWeakCache, destroyPriority,
                std::move(onCreate), std::move(onDestroy)
        );
    }
}
#define REG_CONCAT_IMPL(a, b) a##b
#define REG_CONCAT(a, b) REG_CONCAT_IMPL(a,b)

// Usage： AUTO_REGISTER_FACTORY(MyService, "default", [](){ return std::make_shared<MyService>(42); });
#define AUTO_REGISTER_FACTORY(Type, Name, FactoryExpr, ...) \
    static auto REG_CONCAT(_auto_factory_registrar_, __COUNTER__) = \
        registry::makeAutoFactoryRegistrar<Type>(Name, FactoryExpr, ##__VA_ARGS__)

// Usage： AUTO_REGISTER_INSTANCE(MyService, "singleton", std::make_shared<MyService>(99));
#define AUTO_REGISTER_INSTANCE(Type, Name, InstancePtr, ...) \
    static auto REG_CONCAT(_auto_instance_registrar_, __COUNTER__) = \
        registry::AutoInstanceRegistrar<Type>(Name, InstancePtr, ##__VA_ARGS__)

#define SINGLETON